import numpy as np
import torch
from torch import nn
import colorsys
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.metrics import r2_score

class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.input_hidden = nn.Linear(input_size, hidden_size)
        self.hidden_output = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        hidden = torch.sigmoid(self.input_hidden(x))
        output = self.hidden_output(hidden)
        return output

# 加载数据集
def load_test_data(filename):
    data = pd.read_csv(filename)
    X = data.iloc[:, 1:3].values  # 第二列和第三列为输入特征
    H_actual = data.iloc[:, -3].values  # 实际的 H 值
    S_actual = data.iloc[:, -2].values  # 实际的 S 值
    return X, H_actual, S_actual

# 定义预测函数
def predict(x, model):
    x_torch = torch.tensor(x, dtype=torch.float32).to(device)
    output_activations = model(x_torch)
    return output_activations.detach().cpu().numpy()

# 加载训练好的模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_H = torch.load('trained of H_model.pth', map_location=device)
model_H.to(device)
model_H.eval()

model_S = torch.load('trained of S_model.pth', map_location=device)
model_S.to(device)
model_S.eval()

# 加载测试数据集
X_test, H_actual, S_actual = load_test_data('HSB_dataset.txt')

# 预测 H 值
X_test_torch = torch.tensor(X_test, dtype=torch.float32).to(device)
H_predictions = predict(X_test, model_H).flatten()

# 预测 S 值
S_predictions = predict(X_test, model_S).flatten()

# 将 H 值乘以 360
H_predictions = H_predictions * 360
H_actual = H_actual * 360  # 将 H 实际值也乘以 360

# 计算 H 的误差
H_error = np.mean(np.abs(H_predictions - H_actual))
print(f"Prediction error for H: {H_error}")

# 计算 H 的MSE误差
H_MSEerror = np.mean((H_predictions - H_actual) ** 2)
print(f"Prediction MSE error for H: {H_MSEerror}")

# 计算 H 的决定系数（R^2）
H_r2 = r2_score(H_actual, H_predictions)
print(f"R^2 for H: {H_r2}")

# 计算 S 的误差
S_error = np.mean(np.abs(S_predictions - S_actual))
print(f"Prediction error for S: {S_error}")

# 计算 S 的MSE误差
S_MSEerror = np.mean((S_predictions - S_actual) ** 2)
print(f"Prediction MSE error for S: {S_MSEerror}")

# 计算 S 的决定系数（R^2）
S_r2 = r2_score(S_actual, S_predictions)
print(f"R^2 for S: {S_r2}")

# 绘制前100个H值预测值与真实值的对比曲线
plt.figure(figsize=(12, 6))
plt.plot(H_actual[:100], label='Actual Hue Values', color='blue')
plt.plot(H_predictions[:100], label='Predicted Hue Values', color='red', linestyle='--')
plt.title('Comparison of Actual vs Predicted Hue Values (First 100 Samples)')
plt.xlabel('Sample Index')
plt.ylabel('Hue Value')
plt.legend()
plt.grid(True)
plt.show()

# 绘制前100个S值预测值与真实值的对比曲线
plt.figure(figsize=(12, 6))
plt.plot(S_actual[:100], label='Actual Saturation Values', color='blue')
plt.plot(S_predictions[:100], label='Predicted Saturation Values', color='red', linestyle='--')
plt.title('Comparison of Actual vs Predicted Saturation Values (First 100 Samples)')
plt.xlabel('Sample Index')
plt.ylabel('Saturation Value')
plt.legend()
plt.grid(True)
plt.show()

# 设定 B 值为常数1
B_constant = 1

# 将 HSB 转换为 RGB
colors = [np.clip(colorsys.hsv_to_rgb(h / 360, s, B_constant), 0, 1) for h, s in zip(H_predictions, S_predictions)]

# 绘制结果
plt.figure(figsize=(10, 8))
plt.scatter(X_test[:, 1], X_test[:, 0], c=colors, s=30, edgecolors='k', alpha=0.7)
plt.title('Test Dataset with Color Mapping to RGB')
plt.xlabel('Valence')
plt.ylabel('Arousal')
plt.grid(True)
plt.xlim(-1, 1)
plt.ylim(-1, 1)
plt.show()